from langchain import PromptTemplate, LLMChain
from chatglm_loader import ChatGLM

# 选择模型
model = ChatGLM()
model.load_model("C:\\Work\\llm\\ChatGLM2-6B\\THUDM\\chatglm2-6b")

# 构建模板
template = """Question:{question}

Answer:"""
prompt = PromptTemplate(template=template,
                        input_variables=['question'])

# 构建链
llm_chain = LLMChain(prompt=prompt, llm=model)

#调用
# user question
question = "根据以下文本内容回到我的问题：" + "设备类型：东芝压铸机，部位分类：压铸机，故障内容大杠抽插不动作，故障原因有以下几点：1、大杠锁紧信号线被烧坏，2、大杠松开夹紧磁性开关线缆烧毁短路，3、大杠锁紧螺母闭合限位磁性开关被烫坏，4、大杠螺母松开到位磁性开关松动， 5、大杠挡板上附着铝渣，卡滞，6、大杠锁紧螺母内铝屑过多，卡滞，7、大杠抱合螺母电磁阀阀芯卡死，8、大杠插入电磁阀卡滞。" + "设备类型：东芝压铸机，部位分类：压铸机，故障内容安全门报警无动作，故障原因有以下几点：1、安全门滑动轮轴承坏，2、安全门脱落（导向轮固定螺钉松脱），3、门关不到位（行程开关压块位置不当）， 4、 安全门行程开关进水、失效，5、安全门脱落（安全门支架断），6、安全门电机刹车磨损、未打开，7、安全门链条脱落，8、安全门后门信号线被拉断。" +"设备类型：东芝压铸机，部位分类：压铸机，故障内容模具抽芯无动作故障原因有以下几点：1 接线盒内进水，抽芯接头短路，2 定模抽/插芯到位限位开关老化通断不良，3 抽芯到位信号线公母接头接触不良，4 抽芯行程开关接触不良，5 抽芯不动作，循环超时。" + "设备类型：东芝压铸机，部位分类：压铸机，故障内容机床压力异常不增压，故障原因有以下几点：1 液压油不足，2 系统压力不足（主油泵泵芯磨损），3 压射调压阀密封圈磨损， 4 低速和高速手轮开度调整不当，5 高低速切换阀内部卡有铝渣。" + "设备类型：东芝压铸机，部位分类：压铸机，故障内容压射前进动作缓慢，故障原因有以下几点：1 伺服电磁阀卡滞。" + "设备类型：东芝压铸机，部位分类：压铸机，故障内容压射伺服阀关闭限位异常报警，故障原因有以下几点：1 伺服放大板无电源输出。" + "设备类型：东芝压铸机，部位分类：压铸机，故障内容机床润滑油打不起压力，故障原因有以下几点：1 润滑油定量分配器被砸断 2 润滑油泵磨损 3 润滑油管损坏。" +"设备类型：东芝压铸机，部位分类：压铸机，故障内容开模无动作，故障原因有以下几点：1 开合模快慢速调节阀调整不当。" + "我的问题是:安全门没反应的故障原因是什么？"
print(llm_chain.run(question))
#
# with open("./压铸机2.txt", "r", "gbk") as file:
#     question = file.read()
#
# # 非chain构建
# history = []
# current_length = 0
# for answer_result in  model.generatorAnswer(prompt=prompt.format(question=question),
#                       history=history,
#                       streaming=True):
#     resp = answer_result.llm_output["answer"]
#     history = answer_result.history
#     print(resp[current_length:], end="", flush=True)
#     current_length = len(resp)